package com.castle.pdftools.utils;

import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.Iterator;
import javax.imageio.IIOImage;
import javax.imageio.ImageIO;
import javax.imageio.ImageWriteParam;
import javax.imageio.ImageWriter;

import org.apache.pdfbox.cos.COSName;
import org.apache.pdfbox.io.RandomAccessBufferedFileInputStream;
import org.apache.pdfbox.pdfparser.PDFParser;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.pdmodel.PDPageTree;
import org.apache.pdfbox.pdmodel.PDResources;
import org.apache.pdfbox.pdmodel.graphics.PDXObject;
import org.apache.pdfbox.pdmodel.graphics.form.PDFormXObject;
import org.apache.pdfbox.pdmodel.graphics.image.JPEGFactory;
import org.apache.pdfbox.pdmodel.graphics.image.LosslessFactory;
import org.apache.pdfbox.pdmodel.graphics.image.PDImageXObject;

/**
 * PDF压缩工具类
 */
public final class PDFCompressUtils {

    // -- Fields --

    private File input;
    private File output;
    private float compQual = -1;
    private static float compQualDefault = 0f; //
    private boolean tiff = false;


    /**
     * Set the input file.
     *
     * @throws Exception if the file does not exist or
     *                   cannot be read.
     */
    public void setInput(final File f) throws Exception {
        if (f == null || !f.canRead())
            throw new Exception("Can't read input file: " + f == null ? "<null>" : f.toString());
        this.input = f;
    }

    /**
     * Set the output file.
     *
     * @throws Exception if the file does not exist and cannot
     *                   be created or exists and cannot be written to
     */
    public void setOutput(final File f) throws Exception {
        try {
            if (f == null || (!f.createNewFile() && !f.canWrite()))
                throw new Exception("Can't write to output file: " + f == null ? "<null>" : f.toString());
        } catch (IOException e) {
            throw new Exception(e);
        }
        this.output = f;
    }

    // -- Helper methods --


    private PDDocument shrinkMe()
            throws FileNotFoundException, IOException {
        if (compQual < 0)
            compQual = compQualDefault;
        final RandomAccessBufferedFileInputStream rabfis =
                new RandomAccessBufferedFileInputStream(input);
        final PDFParser parser = new PDFParser(rabfis);
        parser.parse();
        final PDDocument doc = parser.getPDDocument();
        final PDPageTree pages = doc.getPages();
        final ImageWriter imgWriter;
        final ImageWriteParam iwp;
        if (tiff) {
            final Iterator<ImageWriter> tiffWriters =
                    ImageIO.getImageWritersBySuffix("png");
            imgWriter = tiffWriters.next();
            iwp = imgWriter.getDefaultWriteParam();
            //iwp.setCompressionMode(ImageWriteParam.MODE_DISABLED);
        } else {
            final Iterator<ImageWriter> jpgWriters =
                    ImageIO.getImageWritersByFormatName("jpeg");
            imgWriter = jpgWriters.next();
            iwp = imgWriter.getDefaultWriteParam();
            iwp.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
            iwp.setCompressionQuality(compQual);
        }
        for (PDPage p : pages) {
            scanResources(p.getResources(), doc, imgWriter, iwp);
        }
        return doc;
    }

    public void runShell() {
        final boolean allThere = !(input == null || output == null);
        if (!allThere) {
            System.out.println("上传文件和输入文件不允许为空");
            return;
        }
        PDDocument doc = null;
        try {
            doc = shrinkMe();
            doc.save(output);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } finally {
            try {
                doc.close();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    private void scanResources(
            final PDResources rList,
            final PDDocument doc,
            final ImageWriter imgWriter,
            final ImageWriteParam iwp)
            throws FileNotFoundException, IOException {
        Iterable<COSName> xNames = rList.getXObjectNames();
        for (COSName xName : xNames) {
            final PDXObject xObj = rList.getXObject(xName);
            if (xObj instanceof PDFormXObject)
                scanResources(((PDFormXObject) xObj).getResources(), doc, imgWriter, iwp);
            if (!(xObj instanceof PDImageXObject))
                continue;
            final PDImageXObject img = (PDImageXObject) xObj;
            System.out.println("Compressing image: " + xName.getName());
            final ByteArrayOutputStream baos = new ByteArrayOutputStream();
            imgWriter.setOutput(ImageIO.createImageOutputStream(baos));
            BufferedImage bi = img.getImage();
            IIOImage iioi;
            if (bi.getTransparency() == BufferedImage.OPAQUE) {
                iioi = new IIOImage(bi, null, null);
            } else if (bi.getTransparency() == BufferedImage.TRANSLUCENT) {
                iioi = new IIOImage(img.getOpaqueImage(), null, null);
            } else {
                iioi = new IIOImage(img.getOpaqueImage(), null, null);
            }
            imgWriter.write(null, iioi, iwp);
            final ByteArrayInputStream bais =
                    new ByteArrayInputStream(baos.toByteArray());
            final PDImageXObject imgNew;
            if (tiff)
                imgNew = LosslessFactory.createFromImage(doc, img.getImage());
            else
                imgNew = JPEGFactory.createFromStream(doc, bais);
            rList.put(xName, imgNew);
        }
    }

}
